#!/usr/bin/env python3
"""
Amit Kohli, Julien Martel, and Anastasios Angelopoulos
August 10, 2020
This script provides a visualization of the ebv-eye data.
"""

import argparse
import os
import matplotlib.pyplot as plt
from image_process.EyeDataset import Frame, color, EyeDataset

parser = argparse.ArgumentParser(description='Arguments for using the eye visualizer')
parser.add_argument('--subject', type=int, default=22, help='which subject to evaluate')
parser.add_argument('--eye', default='left', choices=['left', 'right'],
                    help='Which eye to visualize, left or right')
parser.add_argument('--data_dir', default=os.path.join(os.getcwd(), 'eye_data'),
                    help='absolute path to eye_data/, by default assumes same parent dir as this script')
parser.add_argument('--buffer', type=int, default=1000, help='How many events to store before displaying.')
opt = parser.parse_args()

def display_data(eye_dataset):
    'Displays the data as fast as GUI can render'
    col_buffer = []
    row_buffer = []
    polarity_buffer = []
    s = plt.plot([],[])[0]
    
    init_img_axis = False
    for i, data in enumerate(eye_dataset):
        if type(data) is Frame:
            if not init_img_axis:
                img_axis = plt.imshow(data.img)
                init = True
            else:
                img_axis.set_data(data.img)
            plt.draw()
            plt.pause(0.0001)
        else:
            col_buffer += [data.col]
            row_buffer += [data.row]
            polarity_buffer += [color[data.polarity]]
            if not len(col_buffer) % opt.buffer:
                s.remove()
                s = plt.scatter(col_buffer, row_buffer, color=polarity_buffer, s=1)
                col_buffer.clear()
                row_buffer.clear()
                polarity_buffer.clear()
                plt.pause(0.0001)

    plt.show()


def main():
    eye_dataset = EyeDataset(opt.data_dir, opt.subject)
    if opt.eye == 'left':
        print('Showing the left eye of subject ' + str(opt.subject) + '\n')
        print('Loading Data from ' + opt.data_dir + '..... \n')
        eye_dataset.collect_data(0)
    else:
        print('Showing the right eye of subject ' + str(opt.subject)+ '\n')
        print('Loading Data from ' + opt.data_dir + '..... \n')
        eye_dataset.collect_data(1)

    for i in range(100):
        print(eye_dataset[i])

    # print('Displaying Data using a ' + str(opt.buffer)+ ' event buffer...')
    
    # display_data(eye_dataset)


# python visualize.py --data_dir ./eye_data --subject 9 --eye left --buffer 1000
if __name__ == '__main__':
    main()







